from paddlers.deploy import Predictor
from util import utils
from operator import itemgetter
import numpy as np
from common import config


class paddleInferancer:
    def __init__(self, model_path, use_gpu, post_process, use_windows=False, window_size=None, stride=None):
        '''
        @model_path:paddle静态图模型路径
        @use_gpu:是否启用GPU推理
        @post_process:默认后处理函数 默认后处理函数的参数应为(原始预测标签,输入原图),且返回值为opencv格式或np.ndarray
        @use_windows:是否启用滑窗推理
        @window_size:滑窗推理窗口尺寸
        @stride:滑窗推理滑动步长
        '''
        if use_windows and (window_size == None or stride == None):
            raise Exception("若启用滑窗推理则window_size和stride不为空")
        # print(model_path)
        self.name = model_path.replace("\\","/").split("/")[-1]
        print(self.name)
        self.predictor = Predictor(model_path, use_gpu)
        self.use_windows = use_windows
        self.window_size = window_size
        self.stride = stride
        self.post_process = post_process

    def inferOne(self, im_a, im_b=None):
        if self.use_windows:
            # 启用滑窗推理
            # 把滑窗中的内容全部取出，存在一个列表中
            ori_size = im_a.shape[:2]
            patch_pairs = []
            for rows, cols in utils.WindowGenerator(*ori_size, self.window_size, self.window_size, self.stride, self.stride):
                patch_pairs.append((im_a[rows, cols], im_b[rows, cols]))
            L = np.array(patch_pairs)
            print(L.shape)
            res = self.predictor.predict(patch_pairs)
            prob_patches = map(itemgetter((..., 1)), map(
                itemgetter('score_map'), res))
            prob_map = utils.recons_prob_map(
                prob_patches, ori_size, self.window_size, self.stride)
        else:
            if self.name == "BIT":
                res = self.predictor.predict((im_a, im_b))
            else:
                res = self.predictor.predict(im_a)
            # print(res)
            if self.name == "YOLO":
                prob_map = res
            else:
                prob_map = res['label_map']
        # 推理时长
        time = self.predictor.timer.inference_time_s.value()
        cm_slide, _ = self.post_process(prob_map, im_a)
        return cm_slide, prob_map, time

# 变化检测模型推理器
#CDInferancer = paddleInferancer(config.CD_MODEL_PATH, True, utils.CDpostProcess, True, config.WINDOW_SIZE, config.STRIDE)
CDInferancer = paddleInferancer(config.CD_MODEL_PATH, True, utils.CDpostProcess, False)

# #目标提取模型推理器
SegInferancer = paddleInferancer(config.SEG_MODEL_PATH, True, utils.SegpostProcess, False)

# #目标检测模型推理器
DetInferancer = paddleInferancer(config.Det_MODEL_PATH, True, utils.DetpostProcess, False)

# #地物分类模型推理器

PDetInferancer = paddleInferancer(config.PDet_MODEL_PATH, True, utils.PDetProcess, False)
